package diskv

import "net"
import "fmt"
import "net/rpc"
import "log"
import "time"
import "paxos"
import "sync"
import "sync/atomic"
import "os"
import "syscall"
import "encoding/gob"
import "encoding/base32"
import "math/rand"
import "shardmaster"
import "io/ioutil"
import (
	"bytes"
	"strconv"
)

const Debug = 0

func DPrintf(format string, a ...interface{}) (n int, err error) {
	if Debug > 0 {
		log.Printf(format, a...)
	}
	return
}

type Op struct {
	// Your definitions here.

	Seq   int64
	Op    string
	Key   string
	Value string

	Data   map[string]string
	SeqMap map[int64]bool
	Num    int
}

type DisKV struct {
	mu         sync.Mutex
	l          net.Listener
	me         int
	dead       int32 // for testing
	unreliable int32 // for testing
	sm         *shardmaster.Clerk
	px         *paxos.Paxos
	dir        string // each replica has its own data directory

	gid int64 // my replica group ID

	// Your definitions here.

	data       map[string]string
	seqMap     map[int64]bool
	lastConfig shardmaster.Config
	lastNum    int
}

func (kv *DisKV) save() error {
	var buffer bytes.Buffer
	enc := gob.NewEncoder(&buffer)
	enc.Encode(kv.lastConfig)
	enc.Encode(kv.lastNum)
	enc.Encode(kv.seqMap)
	kv.px.Save(enc)

	if err := ioutil.WriteFile(kv.dir+"/config-temp", buffer.Bytes(), 0666); err != nil {
		return err
	}
	if err := os.Rename(kv.dir+"/config-temp", kv.dir+"/config"); err != nil {
		return err
	}
	for k, v := range kv.data {
		if err := kv.filePut(key2shard(k), k, v); err != nil {
			return err
		}
	}
	return nil
}

func (kv *DisKV) resume() error {
	var buffer bytes.Buffer
	bytes, err := ioutil.ReadFile(kv.dir + "/config")
	if err != nil {
		return err
	}
	buffer.Write(bytes)

	dec := gob.NewDecoder(&buffer)
	dec.Decode(&kv.lastConfig)
	dec.Decode(&kv.lastNum)
	dec.Decode(&kv.seqMap)
	kv.px.Resume(dec)

	for k, v := range kv.lastConfig.Shards {
		if v == kv.gid {
			data := kv.fileReadShard(k)
			for key, value := range data {
				kv.data[key] = value
			}
		}
	}
	return nil
}

//
// these are handy functions that might be useful
// for reading and writing key/value files, and
// for reading and writing entire shards.
// puts the key files for each shard in a separate
// directory.
//

func (kv *DisKV) shardDir(shard int) string {
	d := kv.dir + "/shard-" + strconv.Itoa(shard) + "/"
	// create directory if needed.
	_, err := os.Stat(d)
	if err != nil {
		if err := os.Mkdir(d, 0777); err != nil {
			log.Fatalf("Mkdir(%v): %v", d, err)
		}
	}
	return d
}

// cannot use keys in file names directly, since
// they might contain troublesome characters like /.
// base32-encode the key to get a file name.
// base32 rather than base64 b/c Mac has case-insensitive
// file names.
func (kv *DisKV) encodeKey(key string) string {
	return base32.StdEncoding.EncodeToString([]byte(key))
}

func (kv *DisKV) decodeKey(filename string) (string, error) {
	key, err := base32.StdEncoding.DecodeString(filename)
	return string(key), err
}

// read the content of a key's file.
func (kv *DisKV) fileGet(shard int, key string) (string, error) {
	fullname := kv.shardDir(shard) + "/key-" + kv.encodeKey(key)
	content, err := ioutil.ReadFile(fullname)
	return string(content), err
}

// replace the content of a key's file.
// uses rename() to make the replacement atomic with
// respect to crashes.
func (kv *DisKV) filePut(shard int, key string, content string) error {
	fullname := kv.shardDir(shard) + "/key-" + kv.encodeKey(key)
	tempname := kv.shardDir(shard) + "/temp-" + kv.encodeKey(key)
	if err := ioutil.WriteFile(tempname, []byte(content), 0666); err != nil {
		return err
	}
	if err := os.Rename(tempname, fullname); err != nil {
		return err
	}
	return nil
}

// return content of every key file in a given shard.
func (kv *DisKV) fileReadShard(shard int) map[string]string {
	m := map[string]string{}
	d := kv.shardDir(shard)
	files, err := ioutil.ReadDir(d)
	if err != nil {
		log.Fatalf("fileReadShard could not read %v: %v", d, err)
	}
	for _, fi := range files {
		n1 := fi.Name()
		if n1[0:4] == "key-" {
			key, err := kv.decodeKey(n1[4:])
			if err != nil {
				log.Fatalf("fileReadShard bad file name %v: %v", n1, err)
			}
			content, err := kv.fileGet(shard, key)
			if err != nil {
				log.Fatalf("fileReadShard fileGet failed for %v: %v", key, err)
			}
			m[key] = content
		}
	}
	return m
}

// replace an entire shard directory.
func (kv *DisKV) fileReplaceShard(shard int, m map[string]string) {
	d := kv.shardDir(shard)
	os.RemoveAll(d) // remove all existing files from shard.
	for k, v := range m {
		kv.filePut(shard, k, v)
	}
}

func (kv *DisKV) Get(args *GetArgs, reply *GetReply) error {
	// Your code here.
	kv.mu.Lock()
	defer kv.mu.Unlock()

	if !kv.inShard(args.Key) || args.Num > kv.lastConfig.Num {
		reply.Err = ErrWrongGroup
		return nil
	}

	kv.px.MakePropose(Op{Seq: nrand(), Op: "Get", Key: args.Key})
	value, ok := kv.data[args.Key]
	if ok {
		reply.Err = OK
		reply.Value = value
	} else {
		reply.Err = ErrNoKey
	}
	kv.save()
	return nil
}

// RPC handler for client Put and Append requests
func (kv *DisKV) PutAppend(args *PutAppendArgs, reply *PutAppendReply) error {
	// Your code here.

	kv.mu.Lock()
	defer kv.mu.Unlock()

	if !kv.inShard(args.Key) || args.Num > kv.lastConfig.Num {
		reply.Err = ErrWrongGroup
		return nil
	}

	reply.Err = OK
	if _, ok := kv.seqMap[args.Seq]; ok {
		return nil
	}

	kv.px.MakePropose(Op{Seq: args.Seq, Op: args.Op, Key: args.Key, Value: args.Value})
	kv.save()
	return nil
}

func (kv *DisKV) Join(args *JoinArgs, reply *JoinReply) error {

	kv.mu.Lock()
	defer kv.mu.Unlock()

	kv.px.MakePropose(Op{Seq: nrand(), Op: "Join", Num: args.Num})

	reply.SeqMap = make(map[int64]bool)
	for k, v := range kv.seqMap {
		reply.SeqMap[k] = v
	}

	reply.Data = make(map[string]string)
	for k, v := range kv.data {
		if _, ok := args.Shards[key2shard(k)]; ok {
			reply.Data[k] = v
		}
	}
	return nil
}

func (kv *DisKV) Update(args *UpdateArgs, reply *UpdateReply) error {
	kv.mu.Lock()
	defer kv.mu.Unlock()

	kv.px.MakePropose(Op{Seq: nrand(), Op: "Empty"})

	reply.SeqMap = make(map[int64]bool)
	for k, v := range kv.seqMap {
		reply.SeqMap[k] = v
	}

	reply.Data = make(map[string]string)
	for k, v := range kv.data {
		reply.Data[k] = v
	}
	reply.Config = kv.lastConfig
	reply.ConfigNum = kv.lastNum

	var buffer bytes.Buffer
	kv.px.Save(gob.NewEncoder(&buffer))
	reply.Paxos = buffer.Bytes()
	return nil
}

func (kv *DisKV) handle(obj interface{}) {
	op := obj.(Op)

	fmt.Println(kv.me, kv.gid, op.Op, op.Key, kv.px.LastSeq)
	switch op.Op {
	case "Get":
	case "Put":
		kv.data[op.Key] = op.Value
	case "Append":
		kv.data[op.Key] += op.Value
	case "Join":
		if op.Num > kv.lastConfig.Num {
			kv.lastNum = op.Num
		}
	case "Empty":
	}
	kv.seqMap[op.Seq] = true
}

func checkOp(a, b interface{}) bool {
	return a.(Op).Seq == b.(Op).Seq
}

func subSlice(a, b []int) []int {
	r := make([]int, 0)

	for _, v1 := range a {
		in := false
		for _, v2 := range b {
			if v2 == v1 {
				in = true
				break
			}
		}
		if !in {
			r = append(r, v1)
		}
	}
	return r
}

func (kv *DisKV) inShard(key string) bool {
	if kv.lastNum > kv.lastConfig.Num {
		return false
	} else {
		return kv.lastConfig.Shards[key2shard(key)] == kv.gid
	}
}

//
// Ask the shardmaster if there's a new configuration;
// if so, re-configure.
//
func (kv *DisKV) tick() {
	// Your code here.

	kv.mu.Lock()
	kv.px.Update()
	kv.mu.Unlock()

	config := kv.sm.Query(kv.lastConfig.Num + 1)
	if config.Num == kv.lastConfig.Num && kv.lastConfig.Num == kv.lastNum {
		return
	}

	newShardSlice := config.GetShard(kv.gid)
	lastShardSlice := kv.lastConfig.GetShard(kv.gid)
	joinSlice := subSlice(newShardSlice, lastShardSlice)
	joinGid := kv.lastConfig.GetGidSlice(joinSlice, kv.gid)

	kv.mu.Lock()
	defer kv.mu.Unlock()

	args := &JoinArgs{Num: config.Num}
	for _, gid := range joinGid {
		args.Shards = make(map[int]bool)
		for _, v := range joinSlice {
			if kv.lastConfig.Shards[v] == gid {
				args.Shards[v] = true
			}
		}
		success := false
		for !success && !kv.isdead() {
			for _, v := range kv.lastConfig.Groups[gid] {
				var reply JoinReply
				if !call(v, "DisKV.Join", args, &reply) {
					continue
				}
				for k, v := range reply.Data {
					kv.data[k] = v
				}

				for k, v := range reply.SeqMap {
					kv.seqMap[k] = v
				}
				success = true
				break
			}
		}
	}
	kv.lastConfig = config.Clone()
	kv.lastNum = kv.lastConfig.Num
	kv.save()
}

// tell the server to shut itself down.
// please don't change these two functions.
func (kv *DisKV) kill() {
	atomic.StoreInt32(&kv.dead, 1)
	kv.l.Close()
	kv.px.Kill()
}

// call this to find out if the server is dead.
func (kv *DisKV) isdead() bool {
	return atomic.LoadInt32(&kv.dead) != 0
}

// please do not change these two functions.
func (kv *DisKV) Setunreliable(what bool) {
	if what {
		atomic.StoreInt32(&kv.unreliable, 1)
	} else {
		atomic.StoreInt32(&kv.unreliable, 0)
	}
}

func (kv *DisKV) isunreliable() bool {
	return atomic.LoadInt32(&kv.unreliable) != 0
}

//
// Start a shardkv server.
// gid is the ID of the server's replica group.
// shardmasters[] contains the ports of the
//   servers that implement the shardmaster.
// servers[] contains the ports of the servers
//   in this replica group.
// Me is the index of this server in servers[].
// dir is the directory name under which this
//   replica should store all its files.
//   each replica is passed a different directory.
// restart is false the very first time this server
//   is started, and true to indicate a re-start
//   after a crash or after a crash with disk loss.
//
func StartServer(gid int64, shardmasters []string,
	servers []string, me int, dir string, restart bool) *DisKV {

	kv := new(DisKV)
	kv.me = me
	kv.gid = gid
	kv.sm = shardmaster.MakeClerk(shardmasters)
	kv.dir = dir

	// Your initialization code here.
	// Don't call Join().

	// log.SetOutput(ioutil.Discard)

	gob.Register(Op{})

	rpcs := rpc.NewServer()
	rpcs.Register(kv)

	kv.px = paxos.Make(servers, me, rpcs)
	kv.data = make(map[string]string)
	kv.seqMap = make(map[int64]bool)

	kv.px.SetFunc(checkOp, kv.handle)

	// log.SetOutput(os.Stdout)

	os.Remove(servers[me])
	l, e := net.Listen("unix", servers[me])
	if e != nil {
		log.Fatal("listen error: ", e)
	}
	kv.l = l

	if err := kv.resume(); err != nil {
		config := kv.sm.Query(-1)
		go func() {
			for k, v := range servers {
				if k == kv.me {
					continue
				}
				args := &UpdateArgs{}
				var reply UpdateReply
				if call(v, "DisKV.Update", args, &reply) {
					if reply.Config.Num != config.Num {
						continue
					}
					for k, v := range reply.Data {
						kv.data[k] = v
					}

					for k, v := range reply.SeqMap {
						kv.seqMap[k] = v
					}

					kv.lastConfig = reply.Config.Clone()
					kv.lastNum = reply.ConfigNum
					buffer := bytes.NewBuffer(reply.Paxos)
					dec := gob.NewDecoder(buffer)
					kv.px.Resume(dec)
					break
				}
			}
		}()
	}

	// please do not change any of the following code,
	// or do anything to subvert it.

	go func() {
		for kv.isdead() == false {
			conn, err := kv.l.Accept()
			if err == nil && kv.isdead() == false {
				if kv.isunreliable() && (rand.Int63()%1000) < 100 {
					// discard the request.
					conn.Close()
				} else if kv.isunreliable() && (rand.Int63()%1000) < 200 {
					// process the request but force discard of reply.
					c1 := conn.(*net.UnixConn)
					f, _ := c1.File()
					err := syscall.Shutdown(int(f.Fd()), syscall.SHUT_WR)
					if err != nil {
						fmt.Printf("shutdown: %v\n", err)
					}
					go rpcs.ServeConn(conn)
				} else {
					go rpcs.ServeConn(conn)
				}
			} else if err == nil {
				conn.Close()
			}
			if err != nil && kv.isdead() == false {
				fmt.Printf("DisKV(%v) accept: %v\n", me, err.Error())
				kv.kill()
			}
		}
	}()


	go func() {
		for kv.isdead() == false {
			kv.tick()
			time.Sleep(250 * time.Millisecond)
		}
	}()

	return kv
}
